# -*- coding: utf-8 -*-

"""
针对于0和1文件夹中的每个视频，获取每帧图像中算法检测和跟踪出的乘客位置。
"""

import os
import json
import oyaml as yaml


def read_json_file(json_file_path):
    """
    读取person_tracker_image文件夹下，每个person id文件夹下的每一个Json文件
    返回该文件所对应视频的相对路径，属于该视频中的哪一帧，以及person在该视频帧图像中的位置
    """
    with open(json_file_path) as json_file:
        data = json.load(json_file)
    video_name = data["source"]["video_name"]
    camera_folder = "c0" if data["source"]["cam"] == "0" else "c1"
    video_path = os.path.join(camera_folder, video_name)
    frame_index = data["source"]["frame_index"]
    box = data["annotation"]["persons"][0]["box"]
    return video_path, frame_index, box


def read_single_person_track_folder(single_person_track_folder, root_folder):
    """
    读取person_tracker_image文件夹下，每个person id文件夹
    返回该文件夹属于哪个视频
    在该视频中，算法检测跟踪到的乘客的位置。
    """
    json_file_paths = [os.path.join(single_person_track_folder, name)
                       for name in os.listdir(single_person_track_folder) if ".json" in name]
    video_path, _, __ = read_json_file(json_file_paths[0])
    frame_person_position = dict()
    for json_file_path in json_file_paths:
        _, frame_index, box = read_json_file(json_file_path)
        frame_person_position[frame_index] = box
    return video_path, frame_person_position


def read_person_tracker_image_folder(track_folder, root_folder):
    result = dict()
    single_person_folders = [os.path.join(track_folder, name)
                             for name in os.listdir(track_folder)
                             if os.path.isdir(os.path.join(track_folder, name))]
    for single_person_folder in single_person_folders:
        video_path, frame_person_position = read_single_person_track_folder(
            single_person_folder, root_folder)
        if video_path not in result:
            result[video_path] = dict()
        if "person_index" not in result[video_path]:
            result[video_path]["person_index"] = [1]
        else:
            pre_person_index = result[video_path]["person_index"][-1]
            result[video_path]["person_index"].append(pre_person_index + 1)
        person_index = result[video_path]["person_index"][-1]
        for frame_index, box in frame_person_position.items():
            if frame_index not in result[video_path]:
                result[video_path][frame_index] = dict()
            result[video_path][frame_index][person_index] = box
    return result


def main():
    root_folder = r"F:\data\bus_videos\59092_00E0B453CE0B\2019_3_18"
    track_folder = os.path.join(root_folder, "information", "person_tracker_image")
    predicted_positions = read_person_tracker_image_folder(
        track_folder, root_folder)

    label_folder = os.path.join(root_folder, "information", "label")
    os.makedirs(label_folder, exist_ok=True)
    yaml_file_path = os.path.join(label_folder, "flow_label.yaml")
    with open(yaml_file_path, "w", encoding="utf-8") as yaml_file:
        yaml.dump(
            {"root": root_folder, "predicted_positions": predicted_positions},
            yaml_file,
            allow_unicode=True,
            default_flow_style=False)


if __name__ == "__main__":
    main()
